#!/usr/bin/python3
# -*- coding: utf-8; tab-width: 4; indent-tabs-mode: t -*-

import os
import re
import sys
import time
import mrget
import functools
import subprocess
import mirrors.plugin
from datetime import datetime


class Main:

    def __init__(self):
        self.ZIM_INFO_DICT = {
            "gutenberg": functools.partial(self._filterKiwixSimple, False),
            "ifixit": functools.partial(self._filterKiwixSimple, False),
#           "mooc": None,                                                           # FIXME: filenames are too messy
            "phet": functools.partial(self._filterKiwixSimple, True),               # files in it have no "_all_"
            "stack_exchange": functools.partial(self._filterKiwixSimple, False),
            "ted": self._filterKiwixTed,
            "videos": self._filterKiwixVideos,
            "vikidia": functools.partial(self._filterKiwixMediaWiki, False),
            "wikibooks": functools.partial(self._filterKiwixMediaWiki, False),
            "wikihow": functools.partial(self._filterKiwixMediaWiki, True),         # files in it have no "_all_"
            "wikinews": functools.partial(self._filterKiwixMediaWiki, False),
            "wikipedia": functools.partial(self._filterKiwixMediaWiki, False),
            "wikiquote": functools.partial(self._filterKiwixMediaWiki, False),
            "wikisource": functools.partial(self._filterKiwixMediaWiki, False),
            "wikiversity": functools.partial(self._filterKiwixMediaWiki, False),
            "wikivoyage": functools.partial(self._filterKiwixMediaWiki, False),
            "wiktionary": functools.partial(self._filterKiwixMediaWiki, False),
        }

        # parameters
        self.id = mirrors.plugin.params["id"]
        self.cfg = mirrors.plugin.params["config"]
        self.dataDir = mirrors.plugin.params["storage-file"]["data-directory"]

        # download source
        self.rsyncUrl = None
        self.fileUrlList = []

    def run(self):
        # self.rsyncUrl
        if True:
            filter_key = lambda x: x["protocol"] == "rsync" and x["role"] != "local-mirror"
            urlList = mrget.target_urls("mirror://kiwix", filter_key=filter_key)
            if len(urlList) == 0:
                print("no appropriate rsync source found", file=sys.stderr)
                sys.exit(1)
            self.rsyncUrl = urlList[0]

        # self.fileUrlList
        if True:
            filter_key = lambda x: x["protocol"] in ["http", "https", "ftp"] and x["role"] != "local-mirror"
            self.fileUrlList = mrget.target_urls("mirror://kiwix", filter_key=filter_key)
            if len(self.fileUrlList) == 0:
                print("no appropriate file source found", file=sys.stderr)
                sys.exit(1)

        # print self.rsyncUrl and self.fileUrlList
        print("Found:")
        print("    rsync source: %s" % (self.rsyncUrl))
        for url in self.fileUrlList:
            print("    file source: %s" % (url))

        # download
        fileListDict = dict()
        for zimName, filterFunc in self.ZIM_INFO_DICT.items():
            print("Synchronizing %s..." % (zimName))

            zimPath = os.path.join("zim", zimName)
            zimLocalDir = os.path.join(self.dataDir, zimPath)

            fileList = self._rsyncGetFileList(zimPath)
            fileList = filterFunc(self.cfg.get(zimName, {}), fileList)

            os.makedirs(zimLocalDir, exist_ok=True)
            self._rsync(zimPath, fileList, zimLocalDir)

            fileListDict[zimLocalDir] = fileList

        # delete
        for zimLocalDir, fileList in fileListDict.items():
            for fn in os.listdir(zimLocalDir):
                if fn not in fileList:
                    os.unlink(os.path.join(zimLocalDir, fn))

    def _rsyncGetFileList(self, zimPath):
        ret = []
        while True:
            try:
                url = os.path.join(self.rsyncUrl, zimPath)
                out = subprocess.check_output(["rsync", "--no-motd", "--list-only", "%s/*" % (url)], universal_newlines=True)    # always re-try
                for line in out.split("\n"):
                    # "drwxr-xr-x          2,380 2021/01/18 12:55:02 ./abc" -> "./abc"
                    m = re.fullmatch(r'\S+ +\S+ +\S+ +\S+ +(.*)', line)
                    if m is None:
                        continue
                    if m.group(1) == ".":
                        continue
                    ret.append(m.group(1))
                return ret
            except subprocess.CalledProcessError as e:
                print("rsync listing files failed and try again: %s" % str(e))
                time.sleep(60)

    def _rsync(self, zimPath, fileList, zimLocalDir):
        # we use "-rlptD" insead of "-a" so that the remote user/group is ignored
        cmdList = ["rsync", "-rlptD", "-z", "-v", "--partial", "-H"]
        for fn in fileList:
            cmdList += ["-f", "+ %s" % (fn)]
        cmdList += ["-f", "- *"]
        cmdList += [os.path.join(self.rsyncUrl, zimPath, "*"), zimLocalDir]
        print(cmdList)

        # do work
        while True:
            try:
                subprocess.check_call(cmdList)      # always re-try
                break
            except subprocess.CalledProcessError as e:
                print("rsync failed and try again: %s" % str(e))
                time.sleep(60)

    def _filterKiwixSimple(self, bTrivial, cfg, fileList):
        # "include-lang", "exclude-lang" in config:
        langIncList, langExcList = _LangFilter.getLangIncExcList(cfg)

        ret = []
        for fn in fileList:
            if not bTrivial:
                # example: XXXXXXXX_ab_all_2020-11.zim
                # we only download "_all_" category files
                m = re.fullmatch(r"\S+_(\S+)_all_\S+\.zim", fn)
            else:
                # example: XXXXXXXX_ab_2020-11.zim
                m = re.fullmatch(r"\S+_(\S+)_\S+\.zim", fn)
            if m is None:
                continue
            if _LangFilter.testByLang(langIncList, langExcList, m.group(1)):        # filter by "include-lang", "exclude-lang"
                continue
            ret.append(fn)

        # filter by date
        ret = _DateFilter.filterByDate(ret)

        return ret

    def _filterKiwixTed(self, cfg, fileList):
        # "include-lang", "exclude-lang" in config:
        langIncList, langExcList = _LangFilter.getLangIncExcList(cfg)

        ret = []
        for fn in fileList:
            # example: ted_en_BLABLABLA_2020-11.zim
            m = re.fullmatch(r"\S+_(\S+)_\S+_\S+\.zim", fn)
            if m is None:
                continue
            if _LangFilter.testByLang(langIncList, langExcList, m.group(1)):        # filter by "include-lang", "exclude-lang"
                continue
            ret.append(fn)

        # filter by date
        ret = _DateFilter.filterByDate(ret)

        return ret

    def _filterKiwixVideos(self, cfg, fileList):
        # "include-lang", "exclude-lang" in config:
        langIncList, langExcList = _LangFilter.getLangIncExcList(cfg)

        ret = []
        for fn in fileList:
            # example (with language):    khan-academy-videos_tr_organik-kimya-fen-bilimleri_2020-12.zim
            # example (without language): voa_learning_english-eim-english-in-a-minute_2021-11.zim
            m = re.fullmatch(r"(\S+)_(\S+)_\S+_\S+\.zim", fn)
            if m is None:
                continue
            if m.group(1) == "voa":
                ret.append(fn)
            else:
                if _LangFilter.testByLang(langIncList, langExcList, m.group(2)):    # filter by "include-lang", "exclude-lang"
                    continue
                ret.append(fn)

        # filter by date
        ret = _DateFilter.filterByDate(ret)

        return ret

    def _filterKiwixMediaWiki(self, bTrivial, cfg, fileList):
        # "file-type" in config:
        fileType = _FileTypeFilter.getFileType(cfg)
        # "include-lang", "exclude-lang" in config:
        langIncList, langExcList = _LangFilter.getLangIncExcList(cfg)

        ret = []
        for fn in fileList:
            if not bTrivial:
                # example: XXXXX_ru_all_maxi_2022-07.zim
                # we only download "_all_" category files
                m = re.fullmatch(r"\S+_(\S+)_all_(%s)_\S+\.zim" % (_FileTypeFilter.getFileTypeRegStr()), fn)
            else:
                # example: XXXXX_ru_maxi_2022-07.zim
                m = re.fullmatch(r"\S+_(\S+)_(%s)_\S+\.zim" % (_FileTypeFilter.getFileTypeRegStr()), fn)
            if m is None:
                continue
            if _FileTypeFilter.testByFileType(fileType, m.group(2)):            # filter by "file-type"
                continue
            if _LangFilter.testByLang(langIncList, langExcList, m.group(1)):    # filter by "include-lang", "exclude-lang"
                continue
            ret.append(fn)

        # filter by date
        ret = _DateFilter.filterByDate(ret)

        return ret


class _LangFilter:

    @staticmethod
    def getLangIncExcList(cfg):
        #   wikipedia_ab_all_maxi_2020-11.zim
        #             ^^

        langIncList = []
        langExcList = []
        if "include-lang" in cfg:
            langIncList = cfg["include-lang"]
        if "exclude-lang" in cfg:
            langExcList = cfg["exclude-lang"]
        if len(langIncList) > 0 and len(langExcList) > 0:
            raise Exception("\"include-lang\" and \"exclude-lang\" can not co-exist in config")
        return langIncList, langExcList

    @staticmethod
    def testByLang(langIncList, langExcList, lang):
        if len(langIncList) > 0:
            if lang not in langIncList:
                return True
        else:
            if lang in langExcList:
                return True
        return False


class _FileTypeFilter:

    _allFileTypes = ["maxi", "mini", "nopic"]

    @classmethod
    def getFileTypeRegStr(cls):
        return "|".join(cls._allFileTypes)

    @classmethod
    def getFileType(cls, cfg):
        #   wikipedia_ab_all_maxi_2020-11.zim, wikipedia_ab_all_mini_2019-02.zim, wikipedia_ab_all_nopic_2020-11.zim
        #                    ^^^^                               ^^^^                               ^^^^^

        fileType = "*"
        if "file-type" in cfg:
            if cfg["file-type"] not in ["*"] + cls._allFileTypes:
                raise Exception("invalid \"file-type\" in config")
            fileType = cfg["file-type"]
        return fileType

    @staticmethod
    def testByFileType(cfgFileType, fileType):
        if cfgFileType != "*" and fileType != cfgFileType:
            return True
        return False


class _DateFilter:

    @staticmethod
    def filterByDate(fileList):
        dateDict = dict()
        ret = []
        for tfn in fileList:
            # fn example: wikipedia_ru_all_maxi_2022-07.zim
            m = re.search(r"(.*)_([0-9-]+)\.zim", tfn)
            if m is None:
                continue

            tKey = m.group(1)
            try:
                tDate = datetime.strptime(m.group(2), "%Y-%m")
            except ValueError:
                tDate = None

            if tKey in dateDict:
                if tDate is not None:
                    if tDate > dateDict[tKey][0]:           # so that we only fetch newest files
                        ret.remove(dateDict[tKey][1])
                        dateDict[tKey] = (tDate, tfn)
                        ret.append(tfn)
                else:
                    ret.append(tfn)                         # files with no date information are also fetched
            else:
                if tDate is not None:
                    dateDict[tKey] = (tDate, tfn)
                ret.append(tfn)

        return ret


###############################################################################

if __name__ == "__main__":
    Main().run()
